{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "handled-prayer",
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "historical-situation",
   "metadata": {},
   "source": [
    "# Setting up equations using Automatic Differentiation and abstract equations\n",
    "This tutorial is meant as an introduction to a new framework for defining and working with (non-linear) equations in PorePy. Specifically, the aim is to develop an approach which:\n",
    "1. Gives a transparent way of specifying non-linear, multiphysics and multi-dimensional equations.\n",
    "2. Speeds up assembly of Jacobian matrices, in particular for geometries with many subdomains.\n",
    "3. Is better suited to combine with general linear and non-linear solvers etc.\n",
    "\n",
    "## Disclaimer\n",
    "The framework, referred to as the \"ad framework' (ad = automatic differentiation) is currently (Spring 2021) under more or less active development. The below tutorial is intended to give an overview of the design and use of the framework, to ease adaptation in new projects. Since the code is under active development, the code will change, hopefully, this tutorial will keep track. In the same spirit, the tutorial will strive to point to uncertainties on how the code will actually function, indicate code that is likely to change, document best practice when opinions on this exist etc."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "decreased-threshold",
   "metadata": {},
   "source": [
    "## Background\n",
    "Over its first few years of existence, PorePy was mainly applied to linear problems; the development was focused more on mixed-dimensionality than on non-linearities. There were two notable exceptions:\n",
    "1. Simulations of viscous fingering in fractured domains, paper [here](https://arxiv.org/abs/1906.10472)\n",
    "2. Simulations of thermo-poromechanics coupled with deformation of fractures, where the latter was modeled as a contact mechanics problem, see for instance [this paper](https://arxiv.org/abs/2008.06289)\n",
    "\n",
    "The two projects both used a Newton-type approach to solving the resulting linear system, but took fundamentally different approaches in the linearization: In the contact-mechanics problem, the Jacobian matrix was formed 'by hand' on a block-matrix level, so, to correctly linearize complex expressions, the user was responsible for applying the chain rule correctly on all terms, on all subdomains. In addition to requiring precision of the user, this approach become somewhat cumbersome on the interface between subdomains were extra classes were had to be implemneted to couple different terms (technically, this has to do with the design of the Assembler object; however, for the purpose of this tutorial there is no need to understand this fully).\n",
    "\n",
    "The non-linear transport problem took a different approach: The project implemented Ad, and thereby removed the painstaking implementation of the Jacobian matrix. To see this works, look first at the [tutorial](https://github.com/pmgbergen/porepy/blob/develop/tutorials/automatic_differentiation.ipynb) on Ad in general, and next on the [tutorial](https://github.com/pmgbergen/porepy/blob/develop/tutorials/compressible_flow_with_automatic_differentiation.ipynb) on how to combine Ad with discretization operations in general."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "further-bride",
   "metadata": {},
   "source": [
    "## Scope of the new Ad framework\n",
    "The new approach to Ad can be seen as an extension of the existing functionality, with the following ambitions:\n",
    "1. For the purpose of writing equations, it should be possible to consider multiple grids simultaneously, with no need for for-loops or similar.\n",
    "2. Instead of the immediate evaluation of residuals and derivatives applied in the existing framework, the new approach should apply delayed evaluation.\n",
    "\n",
    "The first point will both improve code redability, and substantially improve runtimes in cases with many subdomains. The latter allows for advanced linear and non-linear solvers, and possibly for automatic updates of discretizations of non-linear terms, both of which will be explored in the future.\n",
    "\n",
    "## Framework components\n",
    "So far, the framework consists of three types of classes:\n",
    "1. Grid-dependent operators, defined on one or multiple subdomain grids. Examples are:\n",
    "    * divergence and trace operators\n",
    "    * boundary conditions, \n",
    "    * projections between mortar and subdomain grids\n",
    "    * projections between sets of subdomains and subsets.\n",
    "2. Variables. These carry the numerical state of the primary variables, and also values at previous time steps and iteration states.\n",
    "3. Discretization objects. These are mainly shells around standard PorePy discretiation methods.\n",
    "4. Classes needed to turn variables and discretizations into equations, linearize them etc.\n",
    "\n",
    "## Test case: A mixed-dimensional grid.\n",
    "As a test case, we define a mixed-dimensional grid, which we for simplicity let be Cartesian"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "moral-nerve",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import porepy as pp\n",
    "\n",
    "import scipy.sparse.linalg as spla\n",
    "\n",
    "# fractures 1 and 2 cross each other in (3, 3)\n",
    "frac_1 = np.array([[2, 2], [2, 4]])\n",
    "frac_2 = np.array([[2, 5], [3, 3]])\n",
    "# fracture 3 is isolated\n",
    "frac_3 = np.array([[6, 6], [1, 5]])\n",
    "\n",
    "gb = pp.meshing.cart_grid([frac_1, frac_2, frac_3], nx=np.array([7, 7]))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "temporal-booking",
   "metadata": {},
   "source": [
    "Next, we define variables on the subdomains and interfaces. This is done as before:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "bigger-outside",
   "metadata": {},
   "outputs": [],
   "source": [
    "# String representations of the variables.\n",
    "pressure_var = 'pressure'\n",
    "mortar_var = 'mortar_flux'\n",
    "\n",
    "# Loop over all subdomains, define a cell centered variable\n",
    "for _, d in gb:\n",
    "    d[pp.PRIMARY_VARIABLES] = {pressure_var: {'cells': 1}}\n",
    "    \n",
    "# Also loop over interfaces\n",
    "for _, d in gb.edges():\n",
    "    d[pp.PRIMARY_VARIABLES] = {mortar_var: {'cells': 1}}\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "occupied-egyptian",
   "metadata": {},
   "source": [
    "Parameter assignmnet is also done as before, see [this tutorial](https://github.com/pmgbergen/porepy/blob/develop/tutorials/parameter_assignment_assembler_setup.ipynb) for details. Specifically, we will consider a mixed-dimensional flow problem."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "interesting-damages",
   "metadata": {},
   "outputs": [],
   "source": [
    "param_key = 'flow'\n",
    "\n",
    "matrix_perm = 1\n",
    "fracture_perm = 1e2\n",
    "\n",
    "interface_diffusivity = 1e2\n",
    "\n",
    "for g, d in gb:\n",
    "    if g.dim == 2:\n",
    "        perm = pp.SecondOrderTensor(matrix_perm * np.ones(g.num_cells))\n",
    "\n",
    "        # Dirichlet conditions on right and left\n",
    "        left = np.where(np.abs(g.face_centers[0] - gb.bounding_box()[0][0]) < 1e-6)[0]\n",
    "        right = np.where(np.abs(g.face_centers[0] - gb.bounding_box()[1][0]) < 1e-6)[0]\n",
    "        bc_cond = ['dir'] * (left.size + right.size)\n",
    "        bc = pp.BoundaryCondition(g, np.hstack((left, right)), bc_cond)\n",
    "        \n",
    "        bc_val = np.zeros(g.num_faces)\n",
    "        bc_val[left] = 1\n",
    "        specified_data = {'second_order_tensor': perm,\n",
    "                         'bc': bc,\n",
    "                         'bc_values': bc_val}\n",
    "        d = pp.initialize_data(g, d, param_key, specified_data)\n",
    "        \n",
    "    else:\n",
    "        perm = pp.SecondOrderTensor(fracture_perm * np.ones(g.num_cells))\n",
    "        \n",
    "        # No-flow Neumann conditions\n",
    "        bc = pp.BoundaryCondition(g)\n",
    "        bc_val = np.zeros(g.num_faces)\n",
    "        specified_data = {'second_order_tensor': perm,\n",
    "                         'bc': bc,\n",
    "                         'bc_values': bc_val}\n",
    "        d = pp.initialize_data(g, d, param_key, specified_data)\n",
    "        \n",
    "# Initialize data for interfaces as well        \n",
    "for e, d in gb.edges():\n",
    "    mg = d['mortar_grid']\n",
    "    kn = interface_diffusivity * np.ones(mg.num_cells)\n",
    "    pp.initialize_data(mg, d, param_key, {'normal_diffusivity': kn})\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "decimal-dealing",
   "metadata": {},
   "source": [
    "We also give numerical values to the pressure and flux variables, just so that we get more interesting numbers below"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "nearby-corporation",
   "metadata": {},
   "outputs": [],
   "source": [
    "for g, d in gb:\n",
    "    pp.set_state(d)\n",
    "    d[pp.STATE][pressure_var] = np.random.rand(g.num_cells)\n",
    "    \n",
    "for e, d in gb.edges():\n",
    "    pp.set_state(d)\n",
    "    d[pp.STATE][mortar_var] = np.random.rand(d['mortar_grid'].num_cells)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "endangered-operator",
   "metadata": {},
   "source": [
    "### Definition of grid-related operators\n",
    "Now, we are ready to apply the new Ad framework to this mixed-dimensional problem. The key to exploit this efficiently (in terms of both user friendliness and computational speed) is to operate on several grids simultaneously. To that end, we make a list of all subdomain grids, and similarly of all the edges (*not* mortar grids - we need to keep the link to the adjacent subdomains).\n",
    "\n",
    "**NOTE**: The order of the grid in the list is important, as it sets the ordering of variables, discretization object etc. It is recommended to define a list of grids and use this throughout to define variables etc. A list of mortar grids should be made similarly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "amber-argentina",
   "metadata": {},
   "outputs": [],
   "source": [
    "grid_list = [g for g, _ in gb]\n",
    "edge_list = [e for e, _ in gb.edges()]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "improved-plaintiff",
   "metadata": {},
   "source": [
    "Now, we can for instance define a joint divergence operator for all subdomains:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "usual-gardening",
   "metadata": {},
   "outputs": [],
   "source": [
    "div = pp.ad.Divergence(grid_list)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "seven-illness",
   "metadata": {},
   "source": [
    "Note that this is not a matrix, but a special object:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "superior-crossing",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "porepy.numerics.ad.grid_operators.Divergence"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "type(div)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "straight-saskatchewan",
   "metadata": {},
   "source": [
    "We will come back to how to translate div into a numerical expression.\n",
    "\n",
    "We can also define merged projection operators between the subdomain and mortar grids. This can be done either on the whole gb, or on parts of it. The ordering of the grids is important, and frankly not completely clear, but the following seems to work (if you get a warning, disregard it; this will be handled at a later point):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "controlling-prisoner",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/eke001/Dropbox/workspace/python/ppdir/src/porepy/grids/mortar_grid.py:628: UserWarning: Is it really meaningful to ask for signs of a one sided mortar grid?\n",
      "  \"Is it really meaningful to ask for signs of a one sided mortar grid?\"\n"
     ]
    }
   ],
   "source": [
    "mortar_proj = pp.ad.MortarProjections(gb=gb)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ready-march",
   "metadata": {},
   "source": [
    "Critically, the initialization defines a list of grids (and edges), just the same way as we did in the grid list, and, since iterations over the grid bucket items uses a fixed order, we're good.\n",
    "\n",
    "Finally, we will need a representation of boundary conditions:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "extensive-pencil",
   "metadata": {},
   "outputs": [],
   "source": [
    "bound_ad = pp.ad.BoundaryCondition(param_key, grids=grid_list)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "specific-messaging",
   "metadata": {},
   "source": [
    "Again, this is not a numerical boundary condition, but rather a way to access given boundary data.\n",
    "\n",
    "### Mixed-dimensional Ad variables\n",
    "The next step is to define Ad representations of the (mixed-dimensional) variables. For this, we need no less than three different steps (fortunately, we can use these objects for other parts below as well). \n",
    "\n",
    "First, define a degree of freedom manager. For users who have been exposed to the Assembler, this is actually part of that class which has been moved to a separate object, which is responsible for keeping track of local which indices belong to which degrees of freedom:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "dirty-wesley",
   "metadata": {},
   "outputs": [],
   "source": [
    "dof_manager = pp.DofManager(gb)  # note: no pp.ad here"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "potential-count",
   "metadata": {},
   "source": [
    "Next, define an EquationManager. This is a class which may be significantly changed in the months to come, but for the moment, it is responsible for providing Ad  representations of the variables."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "amber-dance",
   "metadata": {},
   "outputs": [],
   "source": [
    "equation_manager = pp.ad.EquationManager(gb, dof_manager)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "gorgeous-grace",
   "metadata": {},
   "source": [
    "Finally, we can define Ad variables"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "black-ethiopia",
   "metadata": {},
   "outputs": [],
   "source": [
    "p = equation_manager.merge_variables([(g, pressure_var) for g in grid_list])\n",
    "lmbda = equation_manager.merge_variables([(e, mortar_var) for e in edge_list])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "center-count",
   "metadata": {},
   "source": [
    "Note that p and lmbda do not have numerical values. What we have done is instead to prepare to:\n",
    "1. Prepare the ground to write equations with the equations\n",
    "2. Prepare for the later translation of the equations to numerical values (values and derivatives)\n",
    "\n",
    "To get some information about the variables, we can type"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "judicial-jacksonville",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Merged variable with name pressure, id 10\n",
      "Composed of 5 variables\n",
      "Degrees of freedom in cells: 1, faces: 0, nodes: 0\n",
      "Total size: 59\n",
      "\n",
      "Merged interface variable with name mortar_flux, id 11\n",
      "Composed of 5 variables\n",
      "Degrees of freedom in cells: 1, faces: 0, nodes: 0\n",
      "Total size: 21\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print(p)\n",
    "print(lmbda)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fatty-guess",
   "metadata": {},
   "source": [
    "### Mixed-dimensional ad equations\n",
    "Next, we turn to discretization. To be compatible with the Ad framework, PorePy discretizations need a wrapper which mainly allows for the delayed evaluation of the expressions. For instance, the Ad version of Mpfa is defined by writing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "ranking-layer",
   "metadata": {},
   "outputs": [],
   "source": [
    "mpfa = pp.ad.MpfaAd(param_key, grid_list)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "destroyed-outdoors",
   "metadata": {},
   "source": [
    "This object, once again, has no numerical values, but is rather an abstract representation of a standard Mpfa discretization. The two versions of Mpfa refer to the discretization matrices resulting from the discretization in similar ways: Mpfa has attributes like 'flux_matrix_key', which specifies where the flux discretization matrix is stored. Similarly, MpfaAd has an attribute 'flux', which, upon parsing of an Ad experession (below), will access the same discretization matrix.\n",
    "\n",
    "To show how this works in action, we can define the flux discretization on subdomain as\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "difficult-edmonton",
   "metadata": {},
   "outputs": [],
   "source": [
    "interior_flux = mpfa.flux * p"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "genetic-stock",
   "metadata": {},
   "source": [
    "In essence, there are two types of Ad objects:\n",
    "1. Atomic objects, like mpfa.flux and p. These can be considered pointers to places in the data dictionary where the numerical values associated with the objects are stored. For instance, p in our example points to a collection of d[pp.STATE][pressure_var], where d is the data dictionary for each of the grids on which p was defined.\n",
    "2. Composite objects, like interior_flux, formed by combining Ad objects (which themselves can be atomic or composites) using basic mathematical operations.\n",
    "\n",
    "These Ad objects are not designed for numerical evaluation by themselves, they can be thought of as recipes for combining discretizations, variables etc (). To parse a recipes, it must be wrapped in an additional layer, termed an Ad Expression, and then provided with a GridBucket, from where it can pull numerical values for variables, discretization matrices and so on."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "defined-daughter",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Ad array of size 134\n",
      "Jacobian is of size (134, 59) and has 174 elements\n"
     ]
    }
   ],
   "source": [
    "eval_flux = pp.ad.Expression(interior_flux, dof_manager)\n",
    "eval_flux.discretize(gb)\n",
    "num_flux = eval_flux.to_ad(gb=gb)\n",
    "print(num_flux)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "reasonable-bachelor",
   "metadata": {},
   "source": [
    "We note that num_flux has the size of the total number of faces in the grids, and that its Jacobian matrix is a mapping from cells to faces.\n",
    "\n",
    "On a technical level (no need to understand this), composite Ad objcets are implemented as a tree structure, where the leaves are atomic Ad objects. Parsing of the expression is done by identification of these leves, and then use standard forward Ad to evaluate the composites.\n",
    "\n",
    "We can define more elaborate combinations of variables. The interior_flux object (side note: Even though we just wrapped it into an Expression, the original composite Ad object is still alive) represents only the part of the flux caused by pressure variations internal to subdomains. To get the full flux, we need to account for boundary conditions from external boundaries, as well as from internal boundaries to domains of lower dimensions.\n",
    "\n",
    "Note that for the time being, we cannot write 'mpfa.bound_flux * (bound_ad + mortar_proj... * lmbba); the parsing of the expressions do not respect parathesis the way it should. To be improved, hopefully."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "balanced-overhead",
   "metadata": {},
   "outputs": [],
   "source": [
    "full_flux = interior_flux + mpfa.bound_flux * bound_ad + mpfa.bound_flux*mortar_proj.mortar_to_primary_int * lmbda"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "sophisticated-sterling",
   "metadata": {},
   "source": [
    "Now, it is interesting to see what happens when the numerical value of full_flux is computed:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "processed-university",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Size of value array: (134,)\n",
      "Size of Jacobian matrix: (134, 80)\n"
     ]
    }
   ],
   "source": [
    "vals = pp.ad.Expression(full_flux, dof_manager).to_ad(gb)\n",
    "print(f'Size of value array: {vals.val.shape}')\n",
    "print(f'Size of Jacobian matrix: {vals.jac.shape}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "auburn-brunswick",
   "metadata": {},
   "source": [
    "Compare the size of the Jacobian matrix with the size of the matrix for int_flux: The number of rows is still equal to the total number of faces in the grid, but the number of columns has increased to also include derivatives with respect to the mortar variables.\n",
    "\n",
    "We can also compute the projection of the mortar fluxes onto the lower-dimensional subdomains, where they are manifested as sources:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "competitive-original",
   "metadata": {},
   "outputs": [],
   "source": [
    "sources_from_mortar = mortar_proj.mortar_to_secondary_int * lmbda"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "nasty-papua",
   "metadata": {},
   "source": [
    "Put together, we now have the full mass conservation equation on all subdomains:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "dimensional-lexington",
   "metadata": {},
   "outputs": [],
   "source": [
    "conservation = div * full_flux + sources_from_mortar"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "right-brunei",
   "metadata": {},
   "source": [
    "We can also define equations for the interface mortars. To that end, we first define the pressure trace on internal boundaries - the most accurate representation of this trace is a bit complex within Mpfa (and Tpfa) methods"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "behind-heater",
   "metadata": {},
   "outputs": [],
   "source": [
    "pressure_trace_from_high = (mortar_proj.primary_to_mortar_avg * mpfa.bound_pressure_cell * p\n",
    "        + mortar_proj.primary_to_mortar_avg * mpfa.bound_pressure_face * mortar_proj.mortar_to_primary_int * lmbda\n",
    "                 )\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "separated-royal",
   "metadata": {},
   "source": [
    "Next, we define a discretization object for the mortar equation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "cleared-multimedia",
   "metadata": {},
   "outputs": [],
   "source": [
    "robin = pp.ad.RobinCouplingAd(param_key, edge_list)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "introductory-buffalo",
   "metadata": {},
   "source": [
    "Now, we can write the Darcy-type equation for the interface flux"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "suburban-vinyl",
   "metadata": {},
   "outputs": [],
   "source": [
    "interface_flux_eq = (robin.mortar_scaling * (pressure_trace_from_high - \n",
    "                                                    mortar_proj.secondary_to_mortar_avg * p)\n",
    "                     + robin.mortar_discr * lmbda)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "characteristic-animal",
   "metadata": {},
   "source": [
    "### Assemble the system of equations\n",
    "Now, we are ready to assemble the full system, formed by the conservation statement and the interface flux equations. Assembly takes two steps:\n",
    "1. Convert the Ad objects into Expressions, preparing for numerical evaluation.\n",
    "2. Feed the Expressions to the EquationManager, thereby join them together into a coupled system. \n",
    "\n",
    "We can do this as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "choice-treat",
   "metadata": {},
   "outputs": [],
   "source": [
    "eqs = [pp.ad.Expression(conservation, dof_manager), pp.ad.Expression(interface_flux_eq, dof_manager)]\n",
    "equation_manager.equations += eqs"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "green-exclusion",
   "metadata": {},
   "source": [
    "The equation_manager can be used to assemble the coupled linear system, much in the same way as an Expression is evaluated. Before that, the discretization matrices must be constructed.\n",
    "\n",
    "**NOTE**: The computed solution has the interpretation of the update to the existing state, that is, the random values we assigned above. The solution must be distributed in an additive manner. \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "answering-denver",
   "metadata": {},
   "outputs": [],
   "source": [
    "# first discretize\n",
    "equation_manager.discretize(gb)\n",
    "# next assemble the equations\n",
    "A, b = equation_manager.assemble_matrix_rhs()\n",
    "\n",
    "# Solve, system, note the minus sign on the right hand side\n",
    "solution = spla.spsolve(A, b)\n",
    "\n",
    "# Distribute variable to local data dictionaries\n",
    "dof_manager.distribute_variable(solution, additive=True)\n",
    "\n",
    "exporter = pp.Exporter(gb, 'ad_test')\n",
    "exporter.write_vtu([pressure_var])\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "about-income",
   "metadata": {},
   "source": [
    "## What have we done\n",
    "We summarize the steps needed to define an equation:\n",
    "1. Define variables \n",
    "2. Define grid-related operators (not strictly necessary, but most often)\n",
    "3. Define discretizations\n",
    "4. Combine into equations, and evaluate.\n",
    "\n",
    "## More advanced usage\n",
    "Below are a few additional techniques which are needed to define other types of equations (to be covered in a more elaborate set of tutorials in the future?):\n",
    "\n",
    "To access the state of a variable on the previous time step, do"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "insured-petroleum",
   "metadata": {},
   "outputs": [],
   "source": [
    "p_prev = p.previous_timestep()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "hydraulic-incident",
   "metadata": {},
   "source": [
    "To use a variable on only a few subdomains, use subdomain projections:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "upper-monkey",
   "metadata": {},
   "outputs": [],
   "source": [
    "g_2d = gb.grids_of_dimension(2)\n",
    "subdomain_proj = pp.ad.SubdomainProjections(grids=g_2d, gb=gb)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "utility-fields",
   "metadata": {},
   "source": [
    "For examples on how to use these, see the BiotContactMechanics model.\n",
    "\n",
    "## Moving parts\n",
    "As hinted to above, the new Ad framework is under active development. There are currently several known issues and shortcomings, including terms in equations that cannot be handled, function names that will change, parameters to functions that should be included or kicked out etc.\n",
    "\n",
    "Below is an attempt at guessing how the main components of the framework will evolve in the future:\n",
    "* The variables, discretizations and Expressions will likely stay more or less as they are, although variable names, functions etc. may be changed.\n",
    "* The EquationManager will likely evolve, if nothing else because it is the least used part of the code.\n",
    "* Better support for constitutive laws etc.\n",
    "* \n",
    "\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
